#include <stdio.h>
#include <assert.h>
#include <string.h>
#include <stdlib.h>
#include <math.h>

#include "data_file.h"  // 为实现方便，将数据直接放进数组中

/**
 *  定义线性回归模型的参数类型
 */
typedef struct 
{
    int paramCount;
    double beta[0];
} LinearRegression;

/**
 * @brief 根据模型参数和数据特征对单条数据进行预测
 * @param model: 线性模型
 * @param x: 数据特征
 * 
 * @return 预测值
 */
double predict(const LinearRegression *model, const double x[])
{
    double ret = model->beta[model->paramCount - 1];
    for (int j = 0; j < model->paramCount - 1; j++) {
        ret += model->beta[j] * x[j];
    }
    return ret;
}

/**
 * @brief 随机梯度下降(SGD)算法实现，即基于单条数据更新模型参数
 * @param model: 线性模型
 * @param x: 某条样本数据的特征
 * @param y_true: 样本x对应的目标值
 * @param lr: 学习率
 * 
 * @return 参数更新前的误差平方损失
 */
double stochastic_gradient_descent(LinearRegression *model, 
    const double x[], const double y_true, double lr)
{
    double y_estimate = predict(model, x);
    double y_err = y_estimate - y_true;
    double loss = 0.5 * y_err * y_err;
    for (int j = 0; j < model->paramCount - 1; j++) {
        model->beta[j] += - lr * y_err * x[j];
    }
    model->beta[model->paramCount - 1] += - lr * y_err;
    return loss;
}

/**
 * @brief 通过随机梯度下降算法进行模型训练
 * @param model: 线性回归模型
 * @param x: 所有的数据特征
 * @param y: 所有的数据目标指
 * @param n: 样本数量
 * @param epochs: 训练轮次
 * @param lr: 学习率
 * 
 * @return 训练结束时，模型的均方误差损失
 */

double train(LinearRegression *model, 
    const double x[][36], const double y[], int n, 
    int epochs, double lr)
{
    double loss_val = 0;
    for (int epoch = 0; epoch < epochs; epoch++) {
        loss_val = 0;
        for (int i = 0; i < n; i++) {
        loss_val += stochastic_gradient_descent(model, x[i], y[i],lr);
        }
        printf("epoch = %d, mse loss = %.4f\n", epoch + 1, loss_val / n);
    }
    
    return loss_val / n;
}


void save_model(const LinearRegression *model, const char *filename)
{
    FILE *file = fopen(filename, "wb");
    assert(file != NULL);
    fwrite(model, sizeof(LinearRegression) + sizeof(double) * model->paramCount, 1, file);
    fclose(file);
}


void load_model(LinearRegression **pmodel, const char *filename) 
{
    FILE *file = fopen(filename, "rb");
    assert(file != NULL);
    fseek(file, 0, SEEK_END);
    size_t nbytes = ftell(file);
    fseek(file, 0, SEEK_SET);

    *pmodel = malloc(nbytes);
    assert(*pmodel != NULL);

    fread(*pmodel, nbytes, 1, file);
    fclose(file);
}

LinearRegression *init_model(size_t n_features)
{
    LinearRegression *model = malloc(sizeof(LinearRegression) + sizeof(double) *(n_features + 1));
    assert(model != NULL);

    model->paramCount = n_features + 1;
    for (int j = 0; j < model->paramCount; j++) {
        model->beta[j] = 1e-7;
    }
    return model;

}

int main(int argc, const char *argv[]) {
    assert(argc == 2);
    if (strcmp(argv[1], "train") == 0) {
        printf("train model\n");
        LinearRegression *pmodel = init_model(n_features);
        double loss = train(pmodel, x_train, y_train, n_train_samples, 100, 1e-3);
        printf("final loss = %.7f\n", loss);
        save_model(pmodel, "linear_model.bin");
        free(pmodel);
    }

    if (strcmp(argv[1], "test") == 0) {
        printf("test model\n");
        LinearRegression *pmodel = NULL;
        load_model(&pmodel, "linear_model.bin");
        double loss = 0;

        for (int i = 0; i < n_test_samples; i++) {
            double y_estimate = predict(pmodel, x_test[i]);
            loss += (y_estimate - y_test[i]) * (y_estimate - y_test[i]);
            printf("第%d条测试数据, 预测收入:%.4f, 实际测收入:%.4f\n", i+1, y_estimate, y_test[i]);
        }
        printf("RMSE = %.4f\n", sqrt(loss/ n_test_samples));
        free(pmodel);
    }

    return 0;
}

